proto[next]: ambient values (clean reimplementation) - #72

Prototype, not for upstream. Supersedes #71, which reached this design through several redesigns and carries their residue. This was written from a specification rather than ported, so the two are directly comparable:

What it does

class Grid(gtx.Container):
    dx: gtx.Static[float]        # folded into the generated code
    nu: gtx.Extern[float]        # passed as a runtime argument


grid = Grid()


@gtx.field_operator
def delta_x(f: IJField) -> IJField:
    return (1.0 / grid.dx) * (f(IDim + 1) - f)     # never a parameter


run(f, out, bind=Grid(dx=0.5, nu=1e-3))

Declarations are annotations bound to plain contextvars.ContextVars. Grid.dx (class access) is the variable, the bind= key; grid.dx (instance access) is its value, so embedded execution sees an ordinary scalar. A filled container binds everything it carries — a grid is one thing semantically and each program picks the parts it reads. Connectivities bind by the same rule with a FieldOffset as the declaration.

A declaration becomes a synthesised program parameter when the program is defined, so from there it is an ordinary argument. The two forms differ only in whether that parameter is listed in static_params: Extern → one compiled program for all values, Static → one per distinct value.

One correction carried over from #71

#71 claimed Static needed no new machinery because gt4py's existing static-argument fold would bake the value in. That was wrong. StaticArg substitution only visits itir_program.body, and an ambient reference lives inside a lowered operator's FunctionDefinition. Verified on both branches by grepping the generated C++:

#71   :  folded=False   (specialised per value, but the literal never substituted)
this  :  folded=True    return [](auto const &f) { return (gtfn::deref(f) * 0.125); };

So Static was paying the recompile cost without the benefit. This branch extends the substitution into the function definitions, filtering the symbol map by each definition's own parameters so a shadowing local cannot be clobbered.

Worth noting the variant-count test could not catch this — it measures specialisation, not folding.

Deliberately out of scope

No content hashing: freeze(), frozen_content_hash() and the hash_offset_provider_items_by_id rewrite from #71 are dropped, and common.py is untouched. That is an orthogonal improvement to a function whose own docstring calls id()-keying unsafe, and it should be proposed on its own.

Also out: ambient fields, connectivities declared inside containers, and a(mesh.V2E) in an operator. FieldOffset stays a module-level declaration and the offset provider is still assembled and passed as today.

Design decisions worth review

• Container subclassing is rejected. Subclasses re-declare, so Base.dx and Derived.dx would be different variables and binding the base would leave the derived unbound. Use composition.
• Partial containers are permitted. Omitting a declaration this program never reads is fine; omitting one it does read raises at call time. Adding a declaration therefore never breaks existing construction sites, at the cost of a typo surfacing later.
• Identity must be stable across restarts. The parameter name reaches the stage fingerprint and hence the build-cache key, so id() is unusable — it would miss the cache every run. A digest of module.qualname.attr is used, and containers that share both (a factory called twice) are rejected with name= as an escape hatch.
• ts.NamespaceType was added instead of the throwaway-class hack in proto[next]: ambient binding of the offset provider #71; keeping the class in ambient.py would have made type_translation import ambient.

Known rough edges

• Grid.dx (class access) is a raw ContextVar, so it owns nothing: bind cannot validate keys or type-check values, and Grid.dx.set(0.5) escapes all scoping. Making it a Declaration that owns the variable would fix all three; noted in the proposal as an open question.
• Two operators closing over different containers under the same variable name are rejected by _get_closure_vars_recursively rather than miscompiled — better than silent shadowing, but not the clean support intended.
• Grid.dx inside an operator body, and grid.dx as an argument in a program body, fail with internal errors rather than DSLErrors.

Tests

tests/next_tests/unit_tests/test_ambient.py (backend-free mechanics) and tests/next_tests/integration_tests/feature_tests/ffront_tests/test_ambient_values.py (backend matrix). Verified independently of the implementation: 331 passed, 4 skipped across both plus test_arg_call_interface.py, CPU backends; pre-commit run clean.